import os
import requests
import m3u8
from concurrent.futures import ThreadPoolExecutor
import subprocess
import shutil

# 配置
M3U8_URL = "https://v30.51cto.com/001c1004fbe271efbb745017e1f80102/47f1913ce7d669b84714a0db48422832-hd-encrypt-stream.m3u8"
SAVE_DIR = "./downloaded_video"
THREADS = 10

# 创建目录
os.makedirs(SAVE_DIR, exist_ok=True)
TS_SEGMENTS_DIR = os.path.join(SAVE_DIR, "ts_segments")
os.makedirs(TS_SEGMENTS_DIR, exist_ok=True)

# 解析 M3U8
def parse_m3u8():
    headers = {
        "User-Agent": "Mozilla/5.0",
        "Referer": "https://www.51cto.com/"
    }
    try:
        response = requests.get(M3U8_URL, headers=headers, timeout=10)
        response.raise_for_status()
        playlist = m3u8.loads(response.text)
    except Exception as e:
        print(f"M3U8 加载失败: {e}")
        return None, []

    key = None
    if playlist.keys and playlist.keys[0]:
        key_url = playlist.keys[0].uri
        base_url = M3U8_URL.rsplit("/", 1)[0]
        key_url = key_url if key_url.startswith("http") else f"{base_url}/{key_url}"
        try:
            key_response = requests.get(key_url, headers=headers, timeout=10)
            key_response.raise_for_status()
            key = (key_url, key_response.content)
        except Exception as e:
            print(f"密钥下载失败: {e}")

    ts_list = [seg.absolute_uri for seg in playlist.segments]
    return key, ts_list

# 下载 TS 分片
def download_ts(ts_url, index, key, iv=None):
    headers = {
        "User-Agent": "Mozilla/5.0",
        "Referer": M3U8_URL
    }
    try:
        response = requests.get(ts_url, headers=headers, timeout=10)
        response.raise_for_status()
        data = response.content

        if key and playlist.keys and playlist.keys[0].iv:
            iv = bytes.fromhex(playlist.keys[0].iv[2:])
            cipher = AES.new(key[1], AES.MODE_CBC, iv=iv)
            data = cipher.decrypt(data)

        filename = os.path.join(TS_SEGMENTS_DIR, f"segment_{index:04d}.ts")
        with open(filename, "wb") as f:
            f.write(data)
        print(f"下载完成 {filename}")
    except Exception as e:
        print(f"下载错误 {ts_url}: {e}")

# 合并 TS 文件
def merge_ts_files(output_mp4):
    ts_files = sorted([
        os.path.join(TS_SEGMENTS_DIR, f)
        for f in os.listdir(TS_SEGMENTS_DIR)
        if f.endswith(".ts")
    ])
    if not ts_files:
        return False

    filelist_path = os.path.join(TS_SEGMENTS_DIR, "filelist.txt")
    with open(filelist_path, "w", encoding="utf-8") as f:
        for ts in ts_files:
            f.write(f"file '{os.path.basename(ts)}'\n")

    command = [
        "ffmpeg",
        "-f", "concat",
        "-safe", "0",
        "-i", filelist_path,
        "-c", "copy",
        output_mp4
    ]
    try:
        subprocess.run(command, check=True)
        print(f"合并完成: {output_mp4}")
        return True
    except Exception as e:
        print(f"FFmpeg 合并失败: {e}")
        return False

# 主流程
def main():
    global playlist  # 用于访问 playlist.keys
    key, ts_list = parse_m3u8()
    if not ts_list:
        return

    print(f"共找到 {len(ts_list)} 个 TS 分片")

    with ThreadPoolExecutor(max_workers=THREADS) as executor:
        for i, ts_url in enumerate(ts_list):
            executor.submit(download_ts, ts_url, i, key)

    output_mp4 = os.path.join(SAVE_DIR, "output.mp4")
    if merge_ts_files(output_mp4):
        try:
            shutil.rmtree(TS_SEGMENTS_DIR)
            print("TS 分片目录已删除")
        except Exception as e:
            print(f"删除 TS 目录失败: {e}")

if __name__ == "__main__":
    main()